package com.hxb.burp.handlers;

import burp.api.montoya.MontoyaApi;
import burp.api.montoya.http.handler.*;
import burp.api.montoya.http.message.requests.HttpRequest;
import burp.api.montoya.logging.Logging;
import com.hxb.burp.utils.AjaxResult;
import com.hxb.burp.utils.HttpUtils;
import com.hxb.burp.utils.IniUtils;
import mjson.Json;

import javax.swing.*;
import java.awt.event.ActionEvent;
import java.awt.event.ActionListener;
import java.awt.event.ComponentEvent;
import java.awt.event.ComponentListener;
import java.nio.charset.StandardCharsets;
import java.util.Arrays;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.stream.Collectors;

/**
 * burp_extendsion2
 *
 * @author 黄小斌
 * @date 2023-10-27
 */
public class WxOfficialHandler implements HttpHandler {
    final private MontoyaApi api;
    final private Logging logger;

    private final ExecutorService threadPool = Executors.newScheduledThreadPool(5);


    private static final String SECTION = "Official";

    public WxOfficialHandler(MontoyaApi api) {
        this.api = api;
        this.logger = api.logging();
        api.userInterface().registerSuiteTab("微信公众号-小程序", defaultPanel());
    }

    public WxOfficialHandler() {
        api = null;
        logger = null;
    }

    JTextField urlField;
    JTextArea consoleArea;
    JScrollPane scrollPane;
    JPanel mainPanel;

    private JPanel defaultPanel() {
        mainPanel = new JPanel();
        mainPanel.setLayout(null);
        mainPanel.setName("haha ");
        JLabel jLabel = new JLabel("服务器地址:");
        jLabel.setBounds(5, 5, 85, 25);
        urlField = new JTextField();
        urlField.setBounds(95, 5, 420, 25);
        JButton testButton = getTestButton();

        JButton saveButton = getSaveButton();
        JButton reSendButton = getResendButton();
        JButton cleanButton = getCleanButton();
        consoleArea = new JTextArea();
        scrollPane = new JScrollPane(consoleArea);
        consoleArea.setEditable(false);
        scrollPane.setVerticalScrollBarPolicy(ScrollPaneConstants.VERTICAL_SCROLLBAR_ALWAYS);
        scrollPane.setAutoscrolls(false);
        scrollPane.setBounds(5, 35, 800, 300);

        mainPanel.add(jLabel);
        mainPanel.add(urlField);
        mainPanel.add(testButton);
        mainPanel.add(saveButton);
        mainPanel.add(reSendButton);
        mainPanel.add(cleanButton);
        mainPanel.add(scrollPane);
        this.initFrame();
        this.resizeListener();

        return mainPanel;
    }


    private void initFrame() {
        urlField.setText(IniUtils.getProperties(SECTION, "serverUrl", "http://localhost:9999/api/wx/session"));
    }

    private String lastSend = "";

    private JButton getTestButton() {
        JButton jButton = new JButton("测试链接");
        jButton.setBounds(520, 5, 90, 25);
        jButton.addActionListener(new ActionListener() {
            @Override
            public void actionPerformed(ActionEvent e) {
                Json test = Json.object();
                test.set("token", "testToken");
                test.set("appId", "testAppId");
                test.set("type", "测试类型");
                sendData(test.toString());
            }
        });
        return jButton;
    }

    private JButton getSaveButton() {
        JButton jButton = new JButton("保存配置");
        jButton.setBounds(615, 5, 90, 25);
        jButton.addActionListener(new ActionListener() {
            @Override
            public void actionPerformed(ActionEvent e) {
                IniUtils.addProperties(SECTION, "serverUrl", urlField.getText());
                printToLog("配置保存成功!");
            }
        });
        return jButton;
    }

    private JButton getResendButton() {
        JButton jButton = new JButton("重新发送");
        jButton.setBounds(710, 5, 90, 25);
        jButton.addActionListener(new ActionListener() {
            @Override
            public void actionPerformed(ActionEvent e) {
                if (lastSend.isEmpty()) {
                    printToLog("最后一次请求数据为空，跳过！");
                    return;
                }
                sendData(lastSend);
            }
        });
        return jButton;
    }

    private JButton getCleanButton() {
        JButton jButton = new JButton("清空日志");
        jButton.setBounds(805, 5, 90, 25);
        jButton.addActionListener(new ActionListener() {
            @Override
            public void actionPerformed(ActionEvent e) {
                consoleArea.setText("");
            }
        });
        return jButton;
    }

    private void printToLog(String message) {
        consoleArea.append(message + "\n");
    }

    private void sendData(String message) {
        this.lastSend = message;
        threadPool.execute(new Runnable() {
            @Override
            public void run() {
                String url = urlField.getText();
                printToLog("请求数据：" + message);
                reDeal(HttpUtils.sendPost(url, message));
            }
        });
    }

    private void resizeListener() {
        if (api != null) {
            int width = this.api.userInterface().swingUtils().suiteFrame().getComponent(0).getWidth() - 10;
            int height = this.api.userInterface().swingUtils().suiteFrame().getComponent(0).getHeight() - 120;
            scrollPane.setBounds(5, 35, width, height);
            api.userInterface().swingUtils().suiteFrame().addComponentListener(new ComponentListener() {
                @Override
                public void componentResized(ComponentEvent e) {
                    int width = mainPanel.getWidth() - 10;
                    int height = mainPanel.getHeight() - 40;
                    scrollPane.setBounds(5, 35, width, height);
                }

                @Override
                public void componentMoved(ComponentEvent e) {

                }

                @Override
                public void componentShown(ComponentEvent e) {

                }

                @Override
                public void componentHidden(ComponentEvent e) {

                }
            });
        }
    }


    private void reDeal(AjaxResult result) {
        if (result.isSuccess()) {
            printToLog("请求结果：" + result.getData().toString());
        } else {
            printToLog("请求异常：" + result.getMsg());
        }
    }

    public static void main(String[] args) {
        WxOfficialHandler handler = new WxOfficialHandler();
        JPanel def = handler.defaultPanel();

        JFrame jFrame = new JFrame("xxxx");
        jFrame.setSize(1000, 500);
        jFrame.add(def);
        jFrame.addComponentListener(new ComponentListener() {
            @Override
            public void componentResized(ComponentEvent e) {
                handler.scrollPane.setBounds(5, 35, def.getWidth() - 10, def.getHeight() - 40);
            }

            @Override
            public void componentMoved(ComponentEvent e) {

            }

            @Override
            public void componentShown(ComponentEvent e) {

            }

            @Override
            public void componentHidden(ComponentEvent e) {

            }
        });
        jFrame.setVisible(true);
    }

    @Override
    public RequestToBeSentAction handleHttpRequestToBeSent(HttpRequestToBeSent requestToBeSent) {

        return RequestToBeSentAction.continueWith(requestToBeSent);
    }

    @Override
    public ResponseReceivedAction handleHttpResponseReceived(HttpResponseReceived responseReceived) {
        HttpRequest request = (HttpRequest) responseReceived.initiatingRequest();
        String requestUrl = request.url();
        if (requestUrl.contains("mp.weixin.qq.com")) {
            try {
                Json json = Json.object();
                String referer = request.header("Referer").toString();
                String token = referer.replaceAll("(.*)(token=)(\\w+)(.*)", "$3");
                String appId = referer.replaceAll("(.*)(appid=)(\\w+)(.*)", "$3");
                json.set("token", token);
                json.set("appId", appId);
                // 获取应用信息
                if (requestUrl.contains("promotion/snsdelivery/snsadvertiser?action=get&token=")) {
                    json.set("type", "appInfo");
                    try {
                        Json body = Json.read(new String(responseReceived.body().getBytes(), StandardCharsets.UTF_8));
                        json.set("appId", body.at("sns_advertiser_info").at("appid").getValue());
                        json.set("nickName", body.at("sns_advertiser_info").at("nick_name").getValue());
                        json.set("headImgUrl", body.at("sns_advertiser_info").at("head_img_url").getValue());
                        sendData(json.toString());
                    } catch (Exception e) {
                        System.err.println(e.getMessage());
                    }
                }
                // 获取Session信息
                if (requestUrl.contains("promotion/publisherplugin/monetize_dashboard?action=roi_trace_v2")) {

                    String[] cookie = request.headerValue("Cookie").split(";");
                    String session = Arrays.stream(cookie).filter(item -> {
                        return item.contains("publisher_session=");
                    }).collect(Collectors.toList()).get(0);
                    json.set("type", "session");
                    json.set("session", session.split("=")[1]);
                    sendData(json.toString());
                }
            } catch (Exception e) {
                System.out.println(e.getMessage());
            }
        }
        return ResponseReceivedAction.continueWith(responseReceived);
    }
}
